package com.caochenlei.tree.trie;

import java.util.HashMap;
import java.util.Map;

//单词查找树的实现类
public class Trie<E> {
    //单词查找树的结点类
    private class Node {
        public Node parent;                     //存储该结点的父结点
        public Map<Character, Node> children;   //存储该结点的子结点
        public Character character;             //存储该结点的字符值
        public E value;                         //存储该结点的对象值
        public boolean isWord;                  //记录是不是单词结尾
    }

    private Node root;                          //记录单词查找树根结点
    private int size;                           //记录单词查找树单词数

    public Trie() {
        this.root = new Node();
        this.root.children = new HashMap<>();
    }

    //获取当前树的根结点
    public Node getRoot() {
        return root;
    }

    //获取当前树的单词数
    public int size() {
        return size;
    }

    //判断当前树是否为空
    public boolean isEmpty() {
        return size == 0;
    }

    //清空当前单词查找树
    public void clear() {
        size = 0;
        root = null;
    }

    //添加、替换操作
    public E add(String key, E value) {
        Node parent = root;
        //循环创建字符路径
        for (int i = 0; i < key.length(); i++) {
            Character character = key.charAt(i);
            Node child = parent.children.get(character);
            if (child == null) {
                child = new Node();
                child.parent = parent;
                child.children = new HashMap<>();
                child.character = character;
                child.value = value;
                child.isWord = false;

                parent.children.put(character, child);
            }
            parent = child;
        }
        //判断单词是否结尾
        if (parent.isWord) {
            E oldValue = parent.value;
            parent.value = value;
            return oldValue;
        }
        //没有结尾则是新增
        parent.isWord = true;
        parent.value = value;
        size++;
        return value;
    }

    //获取指定的结点
    private Node getNode(String key) {
        Node parent = root;
        for (int i = 0; i < key.length(); i++) {
            if (parent == null) return null;
            if (parent.children == null) return null;
            if (parent.children.isEmpty()) return null;
            Character character = key.charAt(i);
            Node child = parent.children.get(character);
            parent = child;
        }
        return parent;
    }

    //根据key获取值
    public E get(String key) {
        Node node = getNode(key);
        return (node != null && node.isWord) ? node.value : null;
    }

    //判断key的存在
    public boolean contains(String key) {
        Node node = getNode(key);
        return (node != null && node.isWord);
    }

    //是否含有前缀串
    public boolean startsWith(String prefix) {
        return getNode(prefix) != null;
    }

    //删除key的操作
    public E remove(String key) {
        //找到最后一个结点
        Node node = getNode(key);
        //如果不是单词结尾
        if (node == null || !node.isWord) {
            return null;
        }
        //开始进行其他操作
        size--;
        E oldValue = node.value;
        //如果还存有子结点
        if (!node.children.isEmpty()) {
            node.isWord = false;
            node.value = null;
            return oldValue;
        }
        //如果不存在子结点
        Node parent;
        while ((parent = node.parent) != null) {
            parent.children.remove(node.character);
            if (parent.isWord || !parent.children.isEmpty()) break;
            node = parent;
        }
        return oldValue;
    }
}
